package com.math.aistudykbr.service.impl;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.math.aistudykbr.mapper.AIMapper;
import com.math.aistudykbr.pojo.entity.MyMessage;
import com.math.aistudykbr.service.AIService;
import com.math.aistudykbr.utils.UserContext;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicLong;

@Service
@RequiredArgsConstructor
@Slf4j
public class AIServiceImpl implements AIService {

    private final AIMapper aiMapper;

    private final RedisTemplate<String, Object> redisTemplate;

    private final OpenAiChatModel openAiChatModel;

    // 普通聊天模板
    private static final String NORMAL_CHAT_TEMPLATE = """
    您是一位专业的AI助手，请使用Markdown格式回答用户的问题，确保回答具有清晰的层级结构：
    1. 使用 # 表示一级标题
    2. 使用 ## 表示二级标题
    3. 使用 ### 表示三级标题
    4. 使用 **文字** 表示加粗文本
    
    请保持回答的简洁和有条理性，将重点内容加粗显示。
    """;

    // 添加工具方法检查调用栈
    private boolean isCalledFromQAService() {
        StackTraceElement[] stackTrace = Thread.currentThread().getStackTrace();
        for (StackTraceElement element : stackTrace) {
            if (element.getClassName().contains("QAServiceImpl") &&
                    element.getMethodName().equals("generateQuestions")) {
                return true;
            }
        }
        return false;
    }

    // 判断消息是否是JSON结构（用于过滤上传文档返回的JSON消息）
    private boolean isJsonMessage(String message) {
        String trimmed = message.trim();
        return (trimmed.startsWith("[") && trimmed.endsWith("]")) || 
               (trimmed.startsWith("{") && trimmed.endsWith("}")) ||
               (trimmed.contains("\"contentTxt\"") && trimmed.contains("\"sonNode\""));
    }

    private final AtomicLong messageCounter = new AtomicLong(0);
    /**
     * 存储最初的消息
     * @param messages
     * @param user
     */
    @Override
    public void storeMessages(String messages, String type, Long user) {
        // 存储消息
        aiMapper.storeMessages(messages, type, user);
    }

    /**
     * 获取历史消息
     * @param userId
     * @return
     */
    @Override
    public List<MyMessage> getHistoryMsg(Long userId) {
        // 1. 先在Redis查询上下文消息
        List<MyMessage> historyMsg = (List<MyMessage>) redisTemplate.opsForValue().get("historyMsg_" + userId);

//        if (historyMsg != null)
//            // 存在历史消息上下文
//            return historyMsg;

        // 2. 不存在历史消息上下文，从数据库查询，则为历史消息
        historyMsg = aiMapper.getHistoryMsg(userId);
        
        // 3. 过滤掉JSON结构的消息回合
        if (historyMsg != null && !historyMsg.isEmpty()) {
            List<MyMessage> filteredMessages = new ArrayList<>();
            
            for (int i = 0; i < historyMsg.size(); i++) {
                MyMessage currentMsg = historyMsg.get(i);
                
                // 如果是用户消息，判断下一条是否为JSON结构的AI回复
                if (Objects.equals(currentMsg.getType(), MessageType.USER.getValue()) 
                        && i + 1 < historyMsg.size()) {
                    
                    MyMessage nextMsg = historyMsg.get(i + 1);
                    if (Objects.equals(nextMsg.getType(), MessageType.ASSISTANT.getValue()) 
                            && isJsonMessage(nextMsg.getMsg())) {
                        // 跳过这一对消息（用户上传和AI返回JSON）
                        i++; // 额外跳过下一条
                        continue;
                    }
                }
                
                // 添加非JSON结构的消息
                filteredMessages.add(currentMsg);
            }
            
            return filteredMessages;
        }

        return historyMsg;
    }

    /**
     * 删除上下文
     * @param userId
     */
    @Override
    public void deleteContext(Long userId) {
        redisTemplate.delete("historyMsg_" + userId);
    }

    /**
     * 删除历史消息
     * @param userId
     */
    @Override
    public void deleteHistoryMsg(Long userId) {

        // 1. 先删除上下文
        redisTemplate.delete("historyMsg_" + userId);
        // 2. 删除数据库中的历史消息
        aiMapper.deleteMessages(userId);
    }

    /**
     * AI对话
     * @param userId
     * @param msg
     * @return
     */
    @Override
    public String chatWithAI(Long userId, String msg) throws JsonProcessingException {
        // 1. 从Redis查询上下文消息，默认是自定义类型，然后再转为Message类型
        List<MyMessage> historyMsgList = (List<MyMessage>) redisTemplate.opsForValue().get("historyMsg_" + userId);
        List<Message> meList = new ArrayList<>();

        if (historyMsgList != null) {
            // 存在历史消息上下文
            MessageHistory(historyMsgList, meList);
        }else{
            // 不存在历史上下文，则为空数据
            historyMsgList = new ArrayList<>();
            historyMsgList.add(new MyMessage(MessageType.USER, msg));
        }

        // 如果是第一条消息或者是QA服务调用，添加系统模板提示
        if (historyMsgList.size() <= 1 || isCalledFromQAService()) {
            // 添加普通聊天的系统提示
            meList.add(0, new SystemMessage(NORMAL_CHAT_TEMPLATE));
        }

        meList.add(new UserMessage(msg));

        // 如果是QA服务调用，跳过历史消息处理
        if (isCalledFromQAService()) {
            meList.add(new UserMessage(msg));
            return openAiChatModel.call(new Prompt(meList))
                    .getResult().getOutput().getText();
        }

        // 2.调用AI接口，获取AI回复
        String aiResponse = openAiChatModel.call(new Prompt(meList)).getResult().getOutput().getText();

        // 3. 存储AI回复
        historyMsgList.add(new MyMessage(MessageType.ASSISTANT, aiResponse));

        // 4. 存储上下文消息
        redisTemplate.opsForValue().set("historyMsg_" + userId, historyMsgList);

        // 5. 存储到数据库
        aiMapper.storeMessages(msg, MessageType.USER.getValue(), userId);
        aiMapper.storeMessages(aiResponse, MessageType.ASSISTANT.getValue(), userId);

        return aiResponse;
    }

    /**
     * AI流式对话
     * @param emitter
     * @param msg
     * @throws JsonProcessingException
     */
    @Override
    public void chatWithAIFlux(SseEmitter emitter, String msg) throws JsonProcessingException {
        // 1. 从Redis查询上下文消息
        Long userId = UserContext.getUser();
        List<MyMessage> historyMsgList = (List<MyMessage>) redisTemplate.opsForValue().get("historyMsg_" + userId);
        List<Message> meList = new ArrayList<>();

        if (historyMsgList != null) {
            // 存在历史消息上下文
            MessageHistory(historyMsgList, meList);
        } else {
            // 不存在历史上下文，则为空数据
            historyMsgList = new ArrayList<>();
        }

        // 如果是第一条消息或历史消息为空，添加系统模板提示
        if (historyMsgList.isEmpty() || meList.isEmpty()) {
            // 添加普通聊天的系统提示
            meList.add(0, new SystemMessage(NORMAL_CHAT_TEMPLATE));
        }

        // 添加用户新消息到历史记录
        historyMsgList.add(new MyMessage(MessageType.USER, msg));
        meList.add(new UserMessage(msg));

        // 用于收集完整的AI回复
        StringBuilder aiResponseBuilder = new StringBuilder();

        // 2.调用AI接口，获取AI流式回复
        List<MyMessage> finalHistoryMsgList = historyMsgList;
        List<MyMessage> finalHistoryMsgList1 = historyMsgList;
        log.info("用户 userID: {}, 开始获取AI流式回复",userId);
        openAiChatModel.stream(new Prompt(meList)).subscribe(
                response -> {
                    try {
                        String content = response.getResult().getOutput().getText();
//                        log.info("AI response chunk: {}", content);

                        long currentSeq = messageCounter.getAndIncrement();

                        // 收集AI回复的每个分块
                        aiResponseBuilder.append(content);

                        // 发送SSE事件
                        emitter.send(SseEmitter.event()
                                .data(ResponseEntity.ok().body(content))
                                .id(String.valueOf(currentSeq))
                                .build());
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                },
                error -> {
                    // 错误处理
                    emitter.completeWithError(error);

                    // 即使出错也保存已接收的部分回复
                    String partialResponse = aiResponseBuilder.toString();
                    if (!partialResponse.isEmpty()) {
                        saveConversation(UserContext.getUser(), msg, partialResponse, finalHistoryMsgList);
                    }
                },
                () -> {
                    // 流式传输完成
                    String completeResponse = aiResponseBuilder.toString();
                    if (!completeResponse.isEmpty()) {
                        // 3. 保存完整的聊天记录
                        saveConversation(userId, msg, completeResponse, finalHistoryMsgList1);
                    }
                }
        );
    }

    private void MessageHistory(List<MyMessage> historyMsgList, List<Message> meList) {
        for (MyMessage historyMsg : historyMsgList) {
            if (Objects.equals(historyMsg.getType(), MessageType.USER.getValue())) {
                meList.add(new UserMessage(historyMsg.getMsg()));
            }else if (Objects.equals(historyMsg.getType(), MessageType.ASSISTANT.getValue())){
                meList.add(new AssistantMessage(historyMsg.getMsg()));
            }
        }
    }

    // 辅助方法：保存对话记录
    private void saveConversation(Long userId, String userMsg, String aiResponse, List<MyMessage> historyMsgList) {
        try {
            log.info("保存用户消息：{}", userId);
            // 添加AI回复到历史记录
            historyMsgList.add(new MyMessage(MessageType.ASSISTANT, aiResponse));

            // 4. 存储上下文消息到Redis
            redisTemplate.opsForValue().set("historyMsg_" + userId, historyMsgList);

            // 5. 存储到数据库
            aiMapper.storeMessages(userMsg, MessageType.USER.getValue(), userId);
            aiMapper.storeMessages(aiResponse, MessageType.ASSISTANT.getValue(), userId);
        } catch (Exception e) {
            log.error("保存聊天记录失败", e);
        }
    }
}
